#include "solverBase.hpp"
#include "sweSolver.hpp"
#include "meshTypeFunctions.hpp"
#include "bc_base_all.hpp"
#include "bc_swe_all.hpp"

#include "structMesh_testFunctions.hpp"

#include <catch2/catch_test_macros.hpp>
#include <catch2/catch_approx.hpp>
#include <fstream>

using Catch::Approx;

template<typename MeshType>
void setDryCells( SWESolver<MeshType>& solver ){
	if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
		solver.d().values()(0, solver.d().mesh().cell(1,1).index() ) = 0;
		solver.d().values()(0, solver.d().mesh().cell(2,1).index() ) = 0;
	} else {
		solver.d().values()(0, 4) = 0;
		solver.d().values()(0, 7) = 0;
	}
}

template<typename MeshType>
void setBottomElevation( SWESolver<MeshType>& solver ){
	if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
		solver.z().values()(0, solver.z().mesh().cell(1,1).index() ) = 0.1;
		solver.z().values()(0, solver.z().mesh().cell(2,1).index() ) = 0.3;
	} else {
		solver.z().values()(0, 4) = 0.1;
		solver.z().values()(0, 7) = 0.3;
	}
}

TEST_CASE( "Solver example", "[pureOutput]" ){
	// using MeshType = Mesh;
	using MeshType = UniMesh;
	// using MeshType = RectMesh;
	// using MeshType = StructMesh;

	Index
		nx {3},
		ny {4},
		ng {1};

	auto mesh { uniMeshAs<MeshType>(nx, ny, 1, 1) };


	Time time {10};
	time.writeTimeInterval(1);

	SWESolver<MeshType> solver;

	solver.v() = Field{mesh, 2};
	solver.d() = Field{mesh, 1};
	solver.z() = Field{mesh, 1};
	solver.v().values().colwise() = Array2s{1, 0};
	solver.d().values() = 1;
	solver.z().values() = 0;

	/* create two dry cells */
	setDryCells(solver);

	solver.n() = ConstantField<Uniform, MeshType>{ mesh, ArrayNs<1>{0.025} };
	
	solver.r() = SourceField<Uniform, MeshType>{ mesh, ArrayNs<1>{0} };
	solver.i() = SourceField<Uniform, MeshType>{ mesh, ArrayNs<1>{0} };

	solver.initialise();

	REQUIRE( solver.calculateTimeStep() == Approx(
		solver.settings().maxCourant /
		( std::sqrt(1.0) + std::sqrt(phys::g*1.0) )
	) );
	
	auto getReshaped = [&](){
		if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
			return std::pair{
				solver.d().values().reshaped( nx + 2*ng, ny + 2*ng ),
				solver.v().values().reshaped( 2*(nx+2*ng), ny+2*ng )
			};
		} else {
			return std::pair{
				solver.d().domainValues().reshaped(   nx, ny ),
				solver.v().domainValues().reshaped( 2*nx, ny )
			};
		}
	};

	auto [dReshaped, vReshaped] = getReshaped();

	std::cout << "Initial values: "
		<< "\nd:\n" << dReshaped
		<< "\nv:\n" << vReshaped
		<< "\n";

	int counter {0};
	while ( !time.atEnd() ){
		if (counter++ >= 1)
			break;
		
		if ( time.atWriteTime() )
			std::cout << "Time: " << time.current()
				<< "\nd:\n" << dReshaped
				<< "\nv:\n" << vReshaped
				<< "\n";

		solver.advance( time );
	}

	std::cout << "Time: " << time.current()
		<< "\nd:\n" << dReshaped
		<< "\nv:\n" << vReshaped
		<< "\n";
}

template<typename MeshType>
void calculateTimeStepTest(){
	Index
		nx {3},
		ny {4},
		ng {1};

	auto mesh { uniMeshAs<MeshType>(nx, ny, 1, 1) };

	SWESolver<MeshType> solver;

	solver.settings().useFriction = false;
	solver.settings().useMuscl = false;

	solver.v() = Field{mesh, 2};
	solver.d() = Field{mesh, 1};
	solver.z() = Field{mesh, 1};
	solver.v().values() = 0;
	solver.d().values() = 1;
	solver.z().values() = 0;

	solver.n() = ConstantField<Uniform, MeshType>{ mesh, ArrayNs<1>{0.025} };

	solver.initialise();

	auto test = [&]( Index i, Index j, Array2s val ){
		auto& v = solver.v().values();
		if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
			v.col( mesh.cell( i, j ).index() ) = val;
		} else {
			Index col;
			if ( j < 0 )
				col = nx*ny + i;
			else if ( j >= ny )
				col = nx*(ny+1) + i;
			else if ( i < 0 )
				col = nx*(ny+2) + j;
			else if ( i >= nx )
				col = nx*(ny+2) + ny + j;
			else
				col = j*nx + i;
			v.col(col) = val;
		}

		REQUIRE( solver.calculateTimeStep() == Approx(
			solver.settings().maxCourant /
			( val.matrix().norm()+std::sqrt(phys::g*1.0) )
		) );
	};

	/* the velocities must be tested in ascending order, because the maximum
	 * velocity is the deciding factor for the time step */

	/* bottom boundary */
	test( 0, -1, {0.2, 0.2} );
	test( 1, -1, {0.3, 0.2} );
	test( 2, -1, {0.4, 0.2} );

	/* top boundary */
	test( 0, 4, {0.5, 0.3} );
	test( 1, 4, {0.6, 0.3} );
	test( 2, 4, {0.7, 0.3} );

	/* left boundary */
	test( -1, 0, {0.8, 0.4} );
	test( -1, 1, {0.9, 0.4} );
	test( -1, 2, {1.0, 0.4} );
	test( -1, 3, {1.1, 0.4} );

	/* right boundary */
	test( 3, 0, {1.2, 0.5} );
	test( 3, 1, {1.3, 0.5} );
	test( 3, 2, {1.4, 0.5} );
	test( 3, 3, {1.5, 0.5} );

	/* inner domain */
	test( 0, 0, {1.6, 0.6} );
	test( 1, 0, {1.7, 0.6} );
	test( 2, 0, {1.8, 0.6} );
	test( 0, 1, {1.9, 0.6} );
	test( 1, 1, {2.0, 0.6} );
	test( 2, 1, {2.1, 0.6} );
	test( 0, 2, {2.2, 0.6} );
	test( 1, 2, {2.3, 0.6} );
	test( 2, 2, {2.4, 0.6} );
	test( 0, 3, {2.5, 0.6} );
	test( 1, 3, {2.6, 0.6} );
	test( 2, 3, {2.7, 0.6} );
}

TEST_CASE( "Fully upwind, 1st order reconstruction, edge-wise", "[pureOutput]" ){

	using MeshType = UniMesh;

	Index
		nx {8},
		ny {8},
		ng {2};

	auto mesh { uniMeshAs<MeshType>(nx, ny, 1, 1) };


	Time time {10};
	time.writeTimeInterval(1);

	SWESolver<MeshType> solver {mesh, time};

	solver.v().values().colwise() = Array2s{1, 0};
	solver.d().values() = 1;
	solver.z().values() = 0;

	/* create two dry cells */
	setDryCells(solver);

	solver.settings().useFriction = false;
	solver.settings().useMuscl = false;
	
	solver.settings().fluxScheme =
		SolverTraits<SWESolver<MeshType>>::FluxScheme::FullyUpwind;
	solver.settings().edgeTraversal = EdgeTraversal::ordered;

	solver.initialise();

	auto getReshaped = [&](){
		if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
			return std::pair{
				solver.d().values().reshaped( nx + 2*ng, ny + 2*ng ),
				solver.v().values().reshaped( 2*(nx+2*ng), ny+2*ng )
			};
		} else {
			return std::pair{
				solver.d().domainValues().reshaped(   nx, ny ),
				solver.v().domainValues().reshaped( 2*nx, ny )
			};
		}
	};

	auto [dReshaped, vReshaped] = getReshaped();

	std::cout << "Initial values: "
		<< "\nd:\n" << dReshaped
		<< "\nv:\n" << vReshaped
		<< "\n";

	int counter {0};
	while ( !time.atEnd() ){
		if (counter++ >= 1)
			break;
		
		if ( time.atWriteTime() )
			std::cout << "Time: " << time.current()
				<< "\nd:\n" << dReshaped
				<< "\nv:\n" << vReshaped
				<< "\n";

		solver.advance( time );
	}
	/* values after advancing 1 time step. */
	std::cout << "Time: " << time.current()
	<< "\nd:\n" << dReshaped
	<< "\nv:\n" << vReshaped
	<< "\n";
}
